Goto

Collaborating Authors

 bure-wasserstein manifold







Test of partial effects for Frechet regression on Bures-Wasserstein manifolds

arXiv.org Machine Learning

In many modern applications, positive definite matrices are often used to summarize the marginal covariance structure among sets of variables. Examples include medical imaging (Dryden et al., 2009; Fillard et al., 2007), neuroscience (Friston, 2011; Kong et al., 2020; Hu et al., 2021) and gene coexpression analysis in single cell genomics. A central challenge in these fields is how to perform regression analysis where the covariance matrix serves as the outcome variable in relation to a set of Euclidean covariates and how to test for the association between these matrix and covariates. Several regression approaches for covariance matrix outcomes have been proposed. Chiu et al. (1996) developed a method that models the elements of the logarithm of the covariance matrix as a linear function of the covariates, but this approach requires estimating a large number of parameters. Hoff & Niu (2012) proposed a regression model where the covariance matrix is expressed as a quadratic function of the explanatory variables. Zou et al. (2017) linked the matrix outcome to a linear combination of similarity matrices derived from the covariates and examined the asymptotic properties of different estimators under this framework. Xu & Li (2025) introduced Fr echet regression with covariate matrix as the outcome.


Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent

Neural Information Processing Systems

We study first-order optimization algorithms for computing the barycenter of Gaussian distributions with respect to the optimal transport metric. Although the objective is geodesically non-convex, Riemannian gradient descent empirically converges rapidly, in fact faster than off-the-shelf methods such as Euclidean gradient descent and SDP solvers. This stands in stark contrast to the best-known theoretical results, which depend exponentially on the dimension. In this work, we prove new geodesic convexity results which provide stronger control of the iterates, yielding a dimension-free convergence rate. Our techniques also enable the analysis of two related notions of averaging, the entropically-regularized barycenter and the geometric median, providing the first convergence guarantees for these problems.


Wasserstein F-tests for Fr\'echet regression on Bures-Wasserstein manifolds

arXiv.org Machine Learning

This paper considers the problem of regression analysis with random covariance matrix as outcome and Euclidean covariates in the framework of Fr\'echet regression on the Bures-Wasserstein manifold. Such regression problems have many applications in single cell genomics and neuroscience, where we have covariance matrix measured over a large set of samples. Fr\'echet regression on the Bures-Wasserstein manifold is formulated as estimating the conditional Fr\'echet mean given covariates $x$. A non-asymptotic $\sqrt{n}$-rate of convergence (up to $\log n$ factors) is obtained for our estimator $\hat{Q}_n(x)$ uniformly for $\left\|x\right\| \lesssim \sqrt{\log n}$, which is crucial for deriving the asymptotic null distribution and power of our proposed statistical test for the null hypothesis of no association. In addition, a central limit theorem for the point estimate $\hat{Q}_n(x)$ is obtained, giving insights to a test for covariate effects. The null distribution of the test statistic is shown to converge to a weighted sum of independent chi-squares, which implies that the proposed test has the desired significance level asymptotically. Also, the power performance of the test is demonstrated against a sequence of contiguous alternatives. Simulation results show the accuracy of the asymptotic distributions. The proposed methods are applied to a single cell gene expression data set that shows the change of gene co-expression network as people age.


Averaging on the Bures-Wasserstein manifold: dimension-free convergence of gradient descent

arXiv.org Artificial Intelligence

We study first-order optimization algorithms for computing the barycenter of Gaussian distributions with respect to the optimal transport metric. Although the objective is geodesically non-convex, Riemannian GD empirically converges rapidly, in fact faster than off-the-shelf methods such as Euclidean GD and SDP solvers. This stands in stark contrast to the best-known theoretical results for Riemannian GD, which depend exponentially on the dimension. In this work, we prove new geodesic convexity results which provide stronger control of the iterates, yielding a dimension-free convergence rate. Our techniques also enable the analysis of two related notions of averaging, the entropically-regularized barycenter and the geometric median, providing the first convergence guarantees for Riemannian GD for these problems.


Wasserstein Gradient Flows for Optimizing Gaussian Mixture Policies

arXiv.org Artificial Intelligence

Robots often rely on a repertoire of previously-learned motion policies for performing tasks of diverse complexities. When facing unseen task conditions or when new task requirements arise, robots must adapt their motion policies accordingly. In this context, policy optimization is the \emph{de facto} paradigm to adapt robot policies as a function of task-specific objectives. Most commonly-used motion policies carry particular structures that are often overlooked in policy optimization algorithms. We instead propose to leverage the structure of probabilistic policies by casting the policy optimization as an optimal transport problem. Specifically, we focus on robot motion policies that build on Gaussian mixture models (GMMs) and formulate the policy optimization as a Wassertein gradient flow over the GMMs space. This naturally allows us to constrain the policy updates via the $L^2$-Wasserstein distance between GMMs to enhance the stability of the policy optimization process. Furthermore, we leverage the geometry of the Bures-Wasserstein manifold to optimize the Gaussian distributions of the GMM policy via Riemannian optimization. We evaluate our approach on common robotic settings: Reaching motions, collision-avoidance behaviors, and multi-goal tasks. Our results show that our method outperforms common policy optimization baselines in terms of task success rate and low-variance solutions.